#include "vtkDoubleArray.h"
#include "vtkPointData.h"
#include "vtkXMLUnstructuredGridWriter.h"
#include "output.h"
#include "mesh.h"


OutputData::OutputData(string _name, int _type) :
    name(_name),
    type(_type)
{
  data.clear();
}

void OutputData::addComponent(double *_data)
{
  data.push_back(_data);
}

OutputData::~OutputData()
{}




Output::Output(Mesh *_mesh, SaveCells _cells) :
    mesh(_mesh), cells(_cells)
{
  int i, ind, j;
  vtkIdType *pts = new vtkIdType[10];
  vtkPoints *points = vtkPoints::New();
  ug = vtkUnstructuredGrid::New();
  ug->Allocate();
  
  // insert points
  if (cells & Save3D)
  {
	  for (i=0; i<mesh->getNvolumes(); i++)
	  {
		for (j=0; j<mesh->getNvolumeDofs(); j++)
		{
		  ind = mesh->getVolumeNode(i,j);
		  points->InsertNextPoint(mesh->getCoord(ind, 0).getVal(), mesh->getCoord(ind, 1).getVal(), mesh->getCoord(ind, 2).getVal());
		}
	  }
  }
  if (cells & Save2D)
  {
	  for (i=0; i<mesh->getNfaces(); i++)
	  {
		for (j=0; j<mesh->getNfaceDofs(); j++)
		{
		  ind = mesh->getFaceNode(i,j);
		  points->InsertNextPoint(mesh->getCoord(ind, 0).getVal(), mesh->getCoord(ind, 1).getVal(), mesh->getCoord(ind, 2).getVal());
		}
	  }
  }
  if (cells & Save1D)
  {
	  for (i=0; i<mesh->getNedges(); i++)
	  {
		for (j=0; j<mesh->getNedgeDofs(); j++)
		{
		  ind = mesh->getEdgeNode(i,j);
		  points->InsertNextPoint(mesh->getCoord(ind, 0).getVal(), mesh->getCoord(ind, 1).getVal(), mesh->getCoord(ind, 2).getVal());
		}
	  }
  }
  ug->SetPoints(points);
  points->Delete();
  
  // insert volumes
  if (cells & Save3D)
  {
	  for (i=0; i<mesh->getNvolumes(); i++)
	  {
		if (mesh->getNvolumeDofs() == 10)
		{
		  for (j=0; j<mesh->getNvolumeDofs(); j++)
			pts[j] = i*mesh->getNvolumeDofs() + VTK_QUAD_TETRA_NODE_ORDER[j];
		}
		else
		{
		  for (j=0; j<mesh->getNvolumeDofs(); j++)
			pts[j] = i*mesh->getNvolumeDofs() + j;
		}
		ug->InsertNextCell(mesh->getVtk3Dtype(), vtkIdType(mesh->getNvolumeDofs()), pts);
	  }
  }
  // insert faces
  if (cells & Save2D)
  {
	  for (i=0; i<mesh->getNfaces(); i++)
	  {
		for (j=0; j<mesh->getNfaceDofs(); j++)
		  pts[j] = mesh->getNvolumes()*mesh->getNvolumeDofs() + i*mesh->getNfaceDofs() + j;
		ug->InsertNextCell(mesh->getVtk2Dtype(), vtkIdType(mesh->getNfaceDofs()), pts);
	  }
  }
  // insert edges
  if (cells & Save1D)
  {
	  for (i=0; i<mesh->getNedges(); i++)
	  {
		for (j=0; j<mesh->getNedgeDofs(); j++)
		  pts[j] = mesh->getNvolumes()*mesh->getNvolumeDofs() + mesh->getNfaces()*mesh->getNfaceDofs()
				  +i*mesh->getNedgeDofs() + j;
		ug->InsertNextCell(mesh->getVtk1Dtype(), vtkIdType(mesh->getNedgeDofs()), pts);
	  }
  }

  delete[] pts;
}

Output::~Output()
{
  for (vector<OutputData*>::iterator i=data.begin(); i!=data.end(); i++) delete *i;
  data.clear();
  ug->Delete();
}

int Output::registerNodeData(string _name, int _type)
{
  data.push_back(new OutputData(_name, _type));
  return data.size()-1;
}

int Output::registerAndAllocateNodeData(string _name, int _type, int size)
{
  int id = data.size();
  data.push_back(new OutputData(_name, _type));
  for (int c=0; c<_type; c++)
    data[id]->data.push_back(new double[size]);
  return id;
}

void Output::addComponent(int data_id, double *_data)
{
  data[data_id]->addComponent(_data);
}

int Output::size()
{
  return (cells & Save3D?1:0)*mesh->getNvolumes()*mesh->getNvolumeDofs()
        +(cells & Save2D?1:0)*mesh->getNfaces()*mesh->getNfaceDofs()
        +(cells & Save1D?1:0)*mesh->getNedges()*mesh->getNedgeDofs();
}



void Output::saveVTK(const char *file)
{
  vtkDoubleArray *da;
  int i, ind=0;
  
  // insert unknowns data
  int N = size();
  for (unsigned int idata=0; idata<data.size(); idata++)
  {
    da = vtkDoubleArray::New();
    da->SetName(data[idata]->name.c_str());
    switch (data[idata]->type)
    {
      case UT_SCALAR:
        da->SetNumberOfComponents(1);
        for (i=0; i<N; i++)
        {
          da->InsertTuple1(i, data[idata]->data[0][i]);
        }
	ind = ug->GetPointData()->AddArray(da);
	ug->GetPointData()->SetActiveAttribute(ind, vtkDataSetAttributes::SCALARS);
        break;
	  
      case UT_VECTOR:
        da->SetNumberOfComponents(3);
        for (i=0; i<N; i++)
        {
          if (data[idata]->data.size() == 2)
          {
            da->InsertTuple3(i, data[idata]->data[0][i], data[idata]->data[1][i], 0);
          }
          else
          {
            da->InsertTuple3(i, data[idata]->data[0][i], data[idata]->data[1][i], data[idata]->data[2][i]);
          }
        }
	ind = ug->GetPointData()->AddArray(da);
	ug->GetPointData()->SetActiveAttribute(ind, vtkDataSetAttributes::VECTORS);
        break;
        
      case UT_SYM_TENSOR:
        da->SetNumberOfComponents(9);
        for (i=0; i<N; i++)
        {
          if (data[idata]->data.size() == 3)
          {
            da->InsertTuple9(i, data[idata]->data[0][i], data[idata]->data[1][i], 0,
                                data[idata]->data[1][i], data[idata]->data[2][i], 0,
                                0, 0, 0);
          }
          else
          {
            da->InsertTuple9(i, data[idata]->data[0][i], data[idata]->data[1][i], data[idata]->data[2][i],
                                data[idata]->data[1][i], data[idata]->data[3][i], data[idata]->data[4][i],
                                data[idata]->data[2][i], data[idata]->data[4][i], data[idata]->data[5][i]);
          }
        }
	ind = ug->GetPointData()->AddArray(da);
	ug->GetPointData()->SetActiveAttribute(ind, vtkDataSetAttributes::TENSORS);
        break;
        
      case UT_TENSOR:
        da->SetNumberOfComponents(9);
        for (i=0; i<N; i++)
        {
          if (data[idata]->data.size() == 4)
          {
            da->InsertTuple9(i, data[idata]->data[0][i], data[idata]->data[1][i], 0,
                                data[idata]->data[2][i], data[idata]->data[3][i], 0,
                                0, 0, 0);
          }
          else
          {
            da->InsertTuple9(i, data[idata]->data[0][i], data[idata]->data[1][i], data[idata]->data[2][i],
                                data[idata]->data[3][i], data[idata]->data[4][i], data[idata]->data[5][i],
                                data[idata]->data[6][i], data[idata]->data[7][i], data[idata]->data[8][i]);
          }
        }
	ind = ug->GetPointData()->AddArray(da);
	ug->GetPointData()->SetActiveAttribute(ind, vtkDataSetAttributes::TENSORS);
        break;
    } // end of unknowns switch
    da->Delete();
  } // end of unknowns loop
  
  // save data to VTK file
  vtkXMLUnstructuredGridWriter *ugw = vtkXMLUnstructuredGridWriter::New();
  ugw->SetInput(ug);
  ugw->SetFileName(file);
  // do not base64() encode the appended data
//  ugw->SetDataModeToAscii();
  ugw->SetEncodeAppendedData(false);
  ugw->Write();
  ugw->Delete();
  
  // remove the data arrays from unstructured grid
  for (i=ug->GetPointData()->GetNumberOfArrays()-1; i>=0; i--)
    ug->GetPointData()->RemoveArray(ug->GetPointData()->GetArrayName(i));
}




void Output::saveGnuplot(const char *fname)
{
  FILE *f;
  int ind;
  
  f = fopen(fname, "w");
  
  for (int e=0; e<mesh->getNvolumes(); e++)
  {
    for (int i=0; i<mesh->getNvolumeDofs(); i++)
    {
      ind = mesh->getVolumeNode(e, i);
      fprintf(f, "%.15e %.15e %.15e ", mesh->getCoord(ind,0).getVal(), mesh->getCoord(ind,1).getVal(), mesh->getCoord(ind,2).getVal());
      for (unsigned int d=0; d<data.size(); d++)
        for (unsigned int c=0; c<data[d]->data.size(); c++)
          fprintf(f, "%.15e ", data[d]->data[c][ind]);
      fprintf(f, "\n");
    }
    fprintf(f, "\n");
  }
  fprintf(f, "\n");
  
  const int face_nodes[] = { 0, 3, 1, 4, 2, 5, 0 };
  CVar val;
  CVec grad;
  for (int e=0; e<mesh->getNfaces(); e++)
  {
    for (int i=0; i<7; i++)
    {
      if (face_nodes[i] >= mesh->getNfaceDofs()) continue;
      ind = mesh->getFaceNode(e, face_nodes[i]);
      fprintf(f, "%.15e %.15e %.15e ", mesh->getCoord(ind,0).getVal(), mesh->getCoord(ind,1).getVal(), mesh->getCoord(ind,2).getVal());
      for (unsigned int d=0; d<data.size(); d++)
        for (unsigned int c=0; c<data[d]->data.size(); c++)
          fprintf(f, "%.15e ", data[d]->data[c][ind]);
      fprintf(f, "\n");
    }
    fprintf(f, "\n\n");
  }
  fprintf(f, "\n");
  
/*  for (int e=0; e<mesh->getNedges(); e++)
  {
    for (int i=0; i<mesh->getNedgeDofs(); i++)
    {
      ind = mesh->getEdgeNode(e, i);
      fprintf(f, "%.15e %.15e %.15e ", mesh->getCoord(ind,0).getVal(), mesh->getCoord(ind,1).getVal(), mesh->getCoord(ind,2).getVal());
      offset=0;
      for (int c=0; c<(int)U.size(); c++)
      {
        fprintf(f, "%.15e ", U[c]->node_value[ind]);
        offset += U[c]->ip1d->getMeshNdof();
      }
      fprintf(f, "\n");
    }
    fprintf(f, "\n");
  }*/
  
  fclose(f);
  
}


